import os
from typing import Dict, List

from langchain import hub
from langchain_community.tools import TavilySearchResults
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain.agents import AgentExecutor,create_openai_tools_agent
from pydantic import BaseModel, Field
from langchain_community.chat_message_histories import ChatMessageHistory,RedisChatMessageHistory
from LimitRedisHistory import LimitedRedisChatMessageHistory
from RedisHistoryWindow import SmartRedisChatMessageHistory
from localModel import LimitRedisHistory

os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_PROJECT'] = 'LLMDEMO'
os.environ['LANGCHAIN_API_KEY'] = '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'
os.environ['LANGSMITH_API_KEY'] = "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"
# 初始化带模板支持的LLM
# pip install redis
llm = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-plus",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)

prompt = hub.pull("hwchase17/openai-functions-agent")
print(prompt.messages)
search = TavilySearchResults(max_results =1)
tools = [search]
agent = create_openai_tools_agent(llm,tools,prompt)
class InMemoryHistory(BaseChatMessageHistory, BaseModel):
    """In memory implementation of chat message history."""
    messages: List[BaseMessage] = Field(default_factory=list)
    def add_messages(self, messages: List[BaseMessage]) -> None:
        """Add a list of messages to the store"""
        super().add_message(messages)

    def clear(self) -> None:
        self.messages = []


store: Dict[str, BaseChatMessageHistory] = {}
REDIS_URL = "redis://:123456@localhost:6379/0"
def get_by_session_id(session_id: str) -> LimitedRedisChatMessageHistory:
    # return RedisChatMessageHistory(session_id, url=REDIS_URL)
    # return LimitedRedisChatMessageHistory(session_id=session_id, url=REDIS_URL,max_messages=2)
    #在原有的RedisChatMessageHistory基础上重写了messages方法，控制最大历史对话轮数，防止传递过多信息给大模型，redis中全量保存用户对话历史
    return SmartRedisChatMessageHistory(session_id=session_id, url=REDIS_URL,chat_memory_window=20)


# def get_by_session_id(session_id: str) -> BaseChatMessageHistory:
#     if session_id not in store:
#         store[session_id] = ChatMessageHistory()
#     return store[session_id]

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
agent_with_history = RunnableWithMessageHistory(
    agent_executor,
    get_by_session_id,
    input_messages_key="input",
    history_messages_key="chat_history"
)
while True:
    input_text = input("请输入问题：")
    if input_text == "exit":
        break
    response = agent_with_history.invoke({"input":input_text},config={"configurable": {"session_id": "foo"}})

